{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 📦 Install Dependencies\n",
    "# \n",
    "# 1. Installs the latest development version of Hugging Face's `diffusers` library directly from GitHub.\n",
    "#    This library is used for working with diffusion models such as Stable Diffusion.\n",
    "!pip install -q -U git+https://github.com/huggingface/diffusers\n",
    "# 2. Installs and upgrades several core Hugging Face and optimization libraries:\n",
    "#    - `transformers`: For using pre-trained language and vision models\n",
    "#    - `accelerate`: Simplifies training across CPUs, GPUs, and distributed setups\n",
    "#    - `wandb`: For logging experiments and visualizing training progress\n",
    "#    - `bitsandbytes`: Enables 8-bit model loading for memory-efficient inference/training\n",
    "#    - `peft`: For Parameter-Efficient Fine-Tuning of large models\n",
    "!pip install -q -U \\\n",
    "    transformers \\\n",
    "    accelerate \\\n",
    "    wandb \\\n",
    "    bitsandbytes \\\n",
    "    peft\n",
    "# 3. Installs additional Python libraries needed for data handling, model input/output,\n",
    "#    and networking or computer vision tasks:\n",
    "#    - `pandas`: For data manipulation and tabular processing\n",
    "#    - `torchvision`: For image transformations and loading datasets (used with PyTorch)\n",
    "#    - `pyarrow`: For efficient I/O and working with Apache Arrow / Parquet formats\n",
    "#    - `sentencepiece`: For subword tokenization used in many NLP models\n",
    "#    - `controlnet_aux`: Adds support functions for ControlNet like HED, Canny, Depth, etc.\n",
    "#    - `scapy`: For packet parsing and crafting, often used in networking/PCAP analysis\n",
    "#    - `gdown`: For downloading files from Google Drive using file IDs\n",
    "#    - `opencv-python`: For computer vision tasks and image manipulation\n",
    "!pip install pandas torchvision pyarrow sentencepiece controlnet_aux scapy gdown opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#As SD3 is gated, before using it with diffusers you first need to go to the Stable Diffusion 3 Medium Hugging Face page, fill in the form and accept the gate. Once you are in, you need to log in so that your system knows you’ve accepted the gate. Use the command below to log in:\n",
    "#Ignore if already logged in.\n",
    "!huggingface-cli login"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the example preprocessed traffic dataset.\n",
    "#\n",
    "# To create your own dataset:\n",
    "# 1) Install nPrint in your environment (see https://nprint.github.io/nprint/).\n",
    "# 2) Collect PCAPs belonging to the same service/application you plan to model.\n",
    "# 3) Convert each PCAP to an nPrint file with the following command:\n",
    "#    nprint -F -1 -P {pcap} -4 -i -6 -t -u -p 0 -c 1024 -W {output_file}\n",
    "# 4) Place all resulting nPrint files in a folder named \"nprint_traffic\".\n",
    "\n",
    "# Download the example publically available preprocessed traffic dataset\n",
    "!gdown --id 1vvneSH0a1WZFPHTafKusOUjNhg7oQioq --output preprocessed_dataset.zip\n",
    "\n",
    "# Unzip and move files into desired directory\n",
    "!unzip -q preprocessed_dataset.zip\n",
    "!mkdir -p nprint_traffic\n",
    "!mv amazon_nprint_traffic/* nprint_traffic/\n",
    "\n",
    "# Clean up\n",
    "!rm -r amazon_nprint_traffic __MACOSX preprocessed_dataset.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------\n",
    "# 🖼️ Convert the nPrint representation of PCAPs into PNG images\n",
    "#\n",
    "# This script is designed to assist in transforming `.nprint` files—\n",
    "# tabular feature representations of network packets—into fixed-size\n",
    "# PNG images that can be used as input for SD fine-tuning.\n",
    "#\n",
    "# 🧾 Input:\n",
    "# - `.nprint` files generated from packet captures (PCAPs) using the\n",
    "#   nPrint tool. Each file is a CSV-like matrix where each row\n",
    "#   represents a single packet and columns correspond to extracted features.\n",
    "#\n",
    "# 🧹 Preprocessing:\n",
    "# - Drops IP address-related columns to avoid injecting identifiable\n",
    "#   or non-generalizable information into the model.\n",
    "# - Maps integer values in the remaining columns to RGBA color tuples\n",
    "#   to visualize numeric features as colored pixels.\n",
    "#\n",
    "# 🧱 Padding:\n",
    "# - Pads each image to a uniform height (default 1024) using a solid\n",
    "#   background to ensure model input consistency across varying packet counts.\n",
    "#\n",
    "# 🎨 Output:\n",
    "# - Saves a PNG file for each `.nprint` file, preserving the packet structure\n",
    "#   as a vertically stacked color-coded image (rows = packets, cols = features).\n",
    "# -----------------------------------------------------------\n",
    "!python ./scripts/nprint_to_png.py -i ./nprint_traffic/ -o ./nprint_traffic_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Compute embeddings:\n",
    "#Generates text prompt embeddings via a Stable Diffusion 3 pipeline and T5 text encoder.\n",
    "#Maps each local PNG image to a unique SHA-256 hash and associates it with the computed embeddings.\n",
    "#Stores the resulting image-hash-to-embedding data in a .parquet file for further processing.\n",
    "#Here we are using the default instance prompt \"pixelated network data for type-0 application traffic\".\n",
    "#But you can configure this. Refer to the compute_embeddings.py script for details on other supported arguments.\n",
    "!python ./scripts/compute_embeddings.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Clear memory\n",
    "import torch\n",
    "import gc\n",
    "\n",
    "def flush():\n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()\n",
    "\n",
    "flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------\n",
    "# 🚀 Train LoRA Adapter on Stable Diffusion 3 (Miniature Setup)\n",
    "#\n",
    "# This command launches training using `accelerate` with DreamBooth-style LoRA tuning,\n",
    "# optimized for quick experimentation or demo runs.\n",
    "#\n",
    "# ⚠️ Current configuration uses:\n",
    "#     - Only 1 training step\n",
    "#     - Small batch size\n",
    "#     - No warmup\n",
    "#     - High learning rate\n",
    "#\n",
    "# Intended **only for testing or verifying training scripts**, NOT for quality results.\n",
    "# For actual training, increase `max_train_steps`, adjust learning rate, and consider\n",
    "# enabling full validation and saving checkpoints.\n",
    "# -----------------------------------------------------------\n",
    "!accelerate launch ./scripts/train_dreambooth_lora_sd3_miniature.py \\\n",
    "  --pretrained_model_name_or_path=\"stabilityai/stable-diffusion-3-medium-diffusers\"  \\\n",
    "  --instance_data_dir=\"nprint_traffic_images\" \\\n",
    "  --data_df_path=\"sample_embeddings.parquet\" \\\n",
    "  --output_dir=\"trained-sd3-lora-miniature\" \\\n",
    "  --mixed_precision=\"fp16\" \\\n",
    "  --instance_prompt=\"pixelated network data for type-0 application traffic\" \\\n",
    "  --train_batch_size=2 \\\n",
    "  --gradient_accumulation_steps=1 --gradient_checkpointing \\\n",
    "  --use_8bit_adam \\\n",
    "  --learning_rate=5e-5 \\\n",
    "  --report_to=\"wandb\" \\\n",
    "  --lr_scheduler=\"constant\" \\\n",
    "  --lr_warmup_steps=0 \\\n",
    "  --max_train_steps=1 \\\n",
    "  --seed=\"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------\n",
    "# 🎯 Inference Pipeline for SD3 + ControlNet + LoRA\n",
    "#\n",
    "# This cell demonstrates the full generation process:\n",
    "#  - Loads Stable Diffusion 3 Medium base model\n",
    "#  - Loads ControlNet (Canny edge-based guidance)\n",
    "#  - Loads LoRA weights fine-tuned on pixelated network traffic\n",
    "#  - Applies edge-conditioned generation using a sample input\n",
    "#\n",
    "# -----------------------------------------------------------\n",
    "flush()\n",
    "import os\n",
    "import torch\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from diffusers import StableDiffusion3ControlNetPipeline, SD3ControlNetModel\n",
    "from diffusers.utils import load_image\n",
    "\n",
    "# Make sure our output folder exists\n",
    "os.makedirs(\"generated_traffic_images\", exist_ok=True)\n",
    "\n",
    "# Base SD 3.0 model\n",
    "base_model_path = \"stabilityai/stable-diffusion-3-medium-diffusers\"\n",
    "\n",
    "# Canny-based ControlNet\n",
    "controlnet_path = \"InstantX/SD3-Controlnet-Canny\"\n",
    "\n",
    "# Load the ControlNet and pipeline\n",
    "controlnet = SD3ControlNetModel.from_pretrained(\n",
    "    controlnet_path, torch_dtype=torch.float16\n",
    ")\n",
    "pipe = StableDiffusion3ControlNetPipeline.from_pretrained(\n",
    "    base_model_path,\n",
    "    controlnet=controlnet,\n",
    ")\n",
    "\n",
    "# Load LoRA weights\n",
    "lora_output_path = \"trained-sd3-lora-miniature\"\n",
    "pipe.load_lora_weights(lora_output_path)\n",
    "\n",
    "# Move pipeline to GPU (half precision)\n",
    "pipe.to(\"cuda\", torch.float16)\n",
    "pipe.enable_sequential_cpu_offload()\n",
    "\n",
    "# ----------------------------------------------------\n",
    "# 1) Convert original control image to Canny edges via OpenCV\n",
    "# ----------------------------------------------------\n",
    "orig_path = \"./scripts/traffic_conditioning_image.png\"\n",
    "orig_bgr = cv2.imread(orig_path)\n",
    "if orig_bgr is None:\n",
    "    raise ValueError(f\"Could not load file: {orig_path}\")\n",
    "\n",
    "# Convert to grayscale\n",
    "gray = cv2.cvtColor(orig_bgr, cv2.COLOR_BGR2GRAY)\n",
    "\n",
    "# Generate Canny edge map (tweak thresholds as needed)\n",
    "edges = cv2.Canny(gray, 100, 200)\n",
    "\n",
    "# Convert single-channel edge map to 3-channel RGB\n",
    "edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)\n",
    "\n",
    "# Convert to PIL for use in ControlNet pipeline\n",
    "control_image = Image.fromarray(edges_rgb)\n",
    "orig_width, orig_height = control_image.size\n",
    "target_width = 1024\n",
    "if orig_width > target_width:\n",
    "    # (left, upper, right, lower)\n",
    "    control_image = control_image.crop((0, 0, target_width, orig_height))\n",
    "    \n",
    "print(\"Displaying Canny control image:\")\n",
    "\n",
    "display(control_image)\n",
    "\n",
    "# ----------------------------------------------------\n",
    "# 3) Set up prompts and run pipeline\n",
    "# ----------------------------------------------------\n",
    "prompt = \"pixelated network data for type-0 application traffic\"\n",
    "generator = torch.manual_seed(0)  # reproducibility\n",
    "\n",
    "# Generate at 1024×1024 to match the new control image\n",
    "image = pipe(\n",
    "    prompt=prompt,\n",
    "    num_inference_steps=20,\n",
    "    generator=generator,\n",
    "    height=1024,\n",
    "    width=1088,\n",
    "    control_image=control_image,\n",
    "    controlnet_conditioning_scale=0.5,  # increase to adhere more strongly to edges\n",
    ").images[0]\n",
    "\n",
    "# ----------------------------------------------------\n",
    "# 4) Save the generated image\n",
    "# ----------------------------------------------------\n",
    "output_path = os.path.join(\"generated_traffic_images\", \"generated_traffic.png\")\n",
    "image.save(output_path)\n",
    "print(f\"Generated image saved to: {output_path}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------\n",
    "# 🔄 Post-Generation Processing Pipeline\n",
    "#\n",
    "# This cell performs a 3-stage transformation of the generated PNG images:\n",
    "#   1. Applies color correction for standardization.\n",
    "#   2. Converts augmented images back into nPrint-compatible feature format.\n",
    "#   3. Applies heuristic corrections and reconstructs valid PCAP files.\n",
    "#\n",
    "# ⚙️ This pipeline enables turning synthetic traffic images\n",
    "#    back into replayable network traffic for evaluation or simulation.\n",
    "# -----------------------------------------------------------\n",
    "\n",
    "# 🎨 Step 1: Color Augmentation\n",
    "# Applies standardized color shifts to improve nprint reconstruction accuracy\n",
    "!python ./scripts/color_processor.py \\\n",
    "  --input_dir=\"./generated_traffic_images\" \\\n",
    "  --output_dir=\"./color_corrected_generated_traffic_images\"\n",
    "# -----------------------------------------------------------\n",
    "# 🔁 Step 2: Image-to-nPrint Conversion\n",
    "# Converts augmented PNG images back into `.nprint` tabular format.\n",
    "#\n",
    "# Uses a reference `.nprint` file to maintain consistent structure and column order.\n",
    "# This step allows diffusion-generated visual traffic to be fed into analysis tools.\n",
    "# -----------------------------------------------------------\n",
    "!python ./scripts/image_to_nprint.py \\\n",
    "  --org_nprint ./scripts/column_example.nprint \\\n",
    "  --input_dir ./color_corrected_generated_traffic_images \\\n",
    "  --output_dir ./generated_nprint\n",
    "# -----------------------------------------------------------\n",
    "# 🧠 Step 3: Heuristic Correction & PCAP Reconstruction\n",
    "#\n",
    "# This step reconstructs a valid and replayable `.pcap` file\n",
    "# from the diffusion-generated `.nprint` representation.\n",
    "# 🔍 Core Functionalities:\n",
    "# ✅ Intra-packet corrections (fixes within individual packets).\n",
    "# 🔁 Inter-packet dependency enforcement.\n",
    "# 🔧 Reconstruction:\n",
    "#   - Save the corrected `.nprint` to disk\n",
    "#   - Call `nprint -W` to convert `.nprint` into `.pcap` using external tool\n",
    "#   - Run Scapy-based checksum updates to ensure IPv4 validity\n",
    "#   - Reconvert final `.pcap` back to `.nprint` (with fixed layout) for downstream tasks\n",
    "# -----------------------------------------------------------\n",
    "!python ./scripts/mass_reconstruction.py \\\n",
    "  --input_dir ./generated_nprint \\\n",
    "  --output_pcap_dir ./replayable_generated_pcaps \\\n",
    "  --output_nprint_dir ./replayable_generated_nprints \\\n",
    "  --formatted_nprint_path ./scripts/correct_format.nprint\n",
    "# Final Pcap is stored in replayable_generated_pcaps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (new_netdiffusion)",
   "language": "python",
   "name": "new_netdiffusion"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
